/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2018-2021. All rights reserved.
 */

package com.baidu.process.cloud.sdk.auth.signer;

import com.baidu.process.cloud.apigateway.sdk.utils.Request;
import com.baidu.process.cloud.sdk.auth.signer.SigningAlgorithm;
import com.baidu.process.cloud.sdk.util.BinaryUtils;
import com.baidu.process.cloud.sdk.util.HttpUtils;

import org.apache.commons.codec.binary.StringUtils;
import org.bouncycastle.crypto.digests.SM3Digest;
import org.openeuler.BGMJCEProvider;

import java.io.UnsupportedEncodingException;
import java.net.URI;
import java.net.URISyntaxException;
import java.nio.charset.StandardCharsets;
import java.security.InvalidKeyException;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.Security;
import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Date;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.SortedMap;
import java.util.TimeZone;
import java.util.TreeMap;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;

public class Signer {
    public static final String LINE_SEPARATOR = "\n";

    public static final String SDK_SIGNING_ALGORITHM = "SDK-HMAC-SHA256";
    public static final String X_SDK_CONTENT_SHA256 = "x-sdk-content-sha256";

    public static final String X_SDK_DATE = "X-Sdk-Date";

    public static final String AUTHORIZATION = "v587sign";

    private static final Pattern AUTHORIZATION_PATTERN_SHA256 = Pattern
            .compile("SDK-HMAC-SHA256\\s+Access=([^,]+),\\s?SignedHeaders=([^,]+),\\s?Signature=(\\w+)");

    private static final Pattern AUTHORIZATION_PATTERN_SM3 = Pattern
            .compile("SDK-HMAC-SM3\\s+Access=([^,]+),\\s?SignedHeaders=([^,]+),\\s?Signature=(\\w+)");

    private static final String LINUX_NEW_LINE = "\n";

    public static final String HOST = "Host";

    public String messageDigestAlgorithm = SDK_SIGNING_ALGORITHM;

    public Signer(String messageDigestAlgorithm) {
        this.messageDigestAlgorithm = messageDigestAlgorithm;
    }

    public Signer() {
    }

    static {
        Security.insertProviderAt(new BGMJCEProvider(), 1);
    }

    public void sign(Request request) throws UnsupportedEncodingException {
        String singerDate = getHeader(request, X_SDK_DATE);
        SimpleDateFormat sdf = new SimpleDateFormat("yyyyMMdd'T'HHmmss'Z'");
        sdf.setTimeZone(TimeZone.getTimeZone("UTC"));

        if (singerDate == null) {
            singerDate = sdf.format(new Date());
            request.addHeader(X_SDK_DATE, singerDate);
        }
        addHostHeader(request);

        String messageDigestContent = calculateContentHash(request);

        String[] signedHeaders = getSignedHeaders(request);

        final String canonicalRequest = createCanonicalRequest(request, signedHeaders, messageDigestContent);

        final byte[] signingKey = deriveSigningKey(request.getSecrect());

        String stringToSign = createStringToSign(canonicalRequest, singerDate);
        byte[] signature = computeSignature(stringToSign, signingKey);
        String signatureResult = buildAuthorizationHeader(signedHeaders, signature, request.getKey());

        request.addHeader(AUTHORIZATION, signatureResult);
    }

    protected String getCanonicalizedResourcePath(String resourcePath) throws UnsupportedEncodingException {
        if (resourcePath == null || resourcePath.isEmpty()) {
            return "/";
        }
        try {
            resourcePath = new URI(resourcePath).getPath();
        } catch (URISyntaxException e) {
            return resourcePath;
        }
        String value = HttpUtils.urlEncode(resourcePath, true);
        if (!value.startsWith("/")) {
            value = "/".concat(value);
        }
        if (!value.endsWith("/")) {
            value = value.concat("/");
        }
        return value;
    }

    protected String getCanonicalizedQueryString(Map<String, List<String>> parameters)
            throws UnsupportedEncodingException {
        final SortedMap<String, List<String>> sorted = new TreeMap<>();

        for (Map.Entry<String, List<String>> entry : parameters.entrySet()) {
            final String encodedParamName = HttpUtils.urlEncode(entry.getKey(), false);
            final List<String> paramValues = entry.getValue();
            final List<String> encodedValues = new ArrayList<>(paramValues.size());
            for (String value : paramValues) {
                encodedValues.add(HttpUtils.urlEncode(value, false));
            }
            Collections.sort(encodedValues);
            sorted.put(encodedParamName, encodedValues);
        }

        final StringBuilder result = new StringBuilder();
        for (Map.Entry<String, List<String>> entry : sorted.entrySet()) {
            for (String value : entry.getValue()) {
                if (result.length() > 0) {
                    result.append("&");
                }
                result.append(entry.getKey()).append("=").append(value);
            }
        }

        return result.toString();
    }

    /**
     * Step 1 of the Signature calculation generate the canonical request.
     *
     * @param request request
     * @param signedHeaders signedHeaders
     * @param messageDigestContent messageDigestContent
     * @return canonicalRequest canonicalRequest
     * @throws UnsupportedEncodingException UnsupportedEncodingException
     */
    protected String createCanonicalRequest(Request request, String[] signedHeaders, String messageDigestContent)
            throws UnsupportedEncodingException {
        return request.getMethod().toString() + LINE_SEPARATOR + getCanonicalizedResourcePath(request.getPath())
                + LINE_SEPARATOR + getCanonicalizedQueryString(request.getQueryStringParams()) + LINE_SEPARATOR
                + getCanonicalizedHeaderString(request, signedHeaders) + LINE_SEPARATOR
                + getSignedHeadersString(signedHeaders) + LINE_SEPARATOR + messageDigestContent;
    }

    /**
     * Step 2 of the Signature calculation
     *
     * @param canonicalRequest canonicalRequest
     * @param singerDate singerDate
     * @return stringToSign stringToSign
     */
    protected String createStringToSign(String canonicalRequest, String singerDate) {
        if (StringUtils.equals(messageDigestAlgorithm, SDK_SIGNING_ALGORITHM)) {
            return messageDigestAlgorithm + LINE_SEPARATOR + singerDate + LINE_SEPARATOR
                    + BinaryUtils.toHex(this.hash(canonicalRequest));
        }
        return messageDigestAlgorithm + LINE_SEPARATOR + singerDate + LINE_SEPARATOR
                + BinaryUtils.toHex(this.hashSm3(canonicalRequest));
    }

    /**
     * Step 3 of the Signature calculation. It involves deriving the signing key and
     * computing the signature
     *
     * @param secret secret
     * @return bytes secret.getBytes
     */
    private byte[] deriveSigningKey(String secret) {
        return secret.getBytes(StandardCharsets.UTF_8);
    }

    protected byte[] sign(byte[] data, byte[] key, SigningAlgorithm algorithm) {
        try {
            Mac mac = Mac.getInstance(algorithm.toString());
            mac.init(new SecretKeySpec(key, algorithm.toString()));
            return mac.doFinal(data);
        } catch (NoSuchAlgorithmException | InvalidKeyException e) {
            return new byte[]{};
        }
    }

    /**
     * Step 3 of the Signature calculation. It involves deriving the signing key and
     * computing the signature
     *
     * @param stringToSign stringToSign
     * @param signingKey signingKey
     * @return bytes bytes
     */
    protected final byte[] computeSignature(String stringToSign, byte[] signingKey) {
        if (StringUtils.equals(messageDigestAlgorithm, SDK_SIGNING_ALGORITHM)) {
            return sign(stringToSign.getBytes(StandardCharsets.UTF_8), signingKey, SigningAlgorithm.HmacSHA256);
        }
        return sign(stringToSign.getBytes(StandardCharsets.UTF_8), signingKey, SigningAlgorithm.HmacSM3);
    }

    /**
     * Creates the authorization header to be included in the request.
     *
     * @param signedHeaders signedHeaders
     * @param signature signature
     * @param accessKey accessKey
     * @return String String
     */
    private String buildAuthorizationHeader(String[] signedHeaders, byte[] signature, String accessKey) {
        final String credential = "Access=" + accessKey;
        final String signerHeaders = "SignedHeaders=" + getSignedHeadersString(signedHeaders);
        final String signatureHeader = "Signature=" + BinaryUtils.toHex(signature);

        return messageDigestAlgorithm + " " + credential + ", " + signerHeaders + ", " + signatureHeader;
    }

    protected String[] getSignedHeaders(Request request) {
        String[] signedHeaders = request.getHeaders().keySet().toArray(new String[0]);
        Arrays.sort(signedHeaders, String.CASE_INSENSITIVE_ORDER);
        return signedHeaders;
    }

    protected String getCanonicalizedHeaderString(Request request, String[] signedHeaders) {
        final Map<String, String> requestHeaders = request.getHeaders();
        StringBuilder buffer = new StringBuilder();
        for (String header : signedHeaders) {
            String key = header.toLowerCase(Locale.getDefault());
            String value = requestHeaders.get(header);
            buffer.append(key).append(":");
            if (value != null) {
                buffer.append(value.trim());
            }

            buffer.append(LINUX_NEW_LINE);
        }
        return buffer.toString();
    }

    protected String getSignedHeadersString(String[] signedHeaders) {
        StringBuilder buffer = new StringBuilder();
        for (String header : signedHeaders) {
            if (buffer.length() > 0) {
                buffer.append(";");
            }
            buffer.append(header.toLowerCase(Locale.getDefault()));
        }

        return buffer.toString();
    }

    protected void addHostHeader(Request request) {
        // requires that we sign the Host header so we
        // have to have it in the request by the time we sign.
        boolean haveHostHeader = false;
        for (String key : request.getHeaders().keySet()) {
            if (HOST.equalsIgnoreCase(key)) {
                haveHostHeader = true;
                break;
            }
        }
        if (!haveHostHeader) {
            request.addHeader(HOST, request.getHost());
        }
    }

    protected String getHeader(Request request, String header) {
        if (header == null) {
            return null;
        }
        Map<String, String> headers = request.getHeaders();
        for (Map.Entry<String, String> entry : headers.entrySet()) {
            if (header.equalsIgnoreCase(entry.getKey())) {
                return entry.getValue();
            }
        }
        return null;
    }

    public boolean verify(Request request) throws UnsupportedEncodingException {
        String singerDate = getHeader(request, X_SDK_DATE);

        String authorization = getHeader(request, AUTHORIZATION);

        Matcher match = AUTHORIZATION_PATTERN_SM3.matcher(authorization);
        if (StringUtils.equals(messageDigestAlgorithm, SDK_SIGNING_ALGORITHM)) {
            match = AUTHORIZATION_PATTERN_SHA256.matcher(authorization);
        }
        if (!match.find()) {
            return false;
        }

        String[] signedHeaders = match.group(2).split(";");
        byte[] signingKey = deriveSigningKey(request.getSecrect());
        String messageDigestContent = calculateContentHash(request);
        String canonicalRequest = createCanonicalRequest(request, signedHeaders, messageDigestContent);
        String stringToSign = createStringToSign(canonicalRequest, singerDate);
        byte[] signature = computeSignature(stringToSign, signingKey);
        String signatureResult = buildAuthorizationHeader(signedHeaders, signature, request.getKey());

        return signatureResult.equals(authorization);
    }

    protected String calculateContentHash(Request request) {
        String content_sha256 = getHeader(request, X_SDK_CONTENT_SHA256);
        if (content_sha256 != null) {
            return content_sha256;
        }
        if (StringUtils.equals(messageDigestAlgorithm, SDK_SIGNING_ALGORITHM)) {
            return BinaryUtils.toHex(hash(request.getBody()));
        }
        return BinaryUtils.toHex(hashSm3(request.getBody()));
    }

    public byte[] hash(String text) {
        try {
            MessageDigest md = MessageDigest.getInstance("SHA-256");
            md.update(text.getBytes(StandardCharsets.UTF_8));
            return md.digest();
        } catch (NoSuchAlgorithmException e) {
            return new byte[]{};
        }
    }

    public byte[] hashSm3(String text) {
        byte[] srcData = text.getBytes(StandardCharsets.UTF_8);
        SM3Digest digest = new SM3Digest();
        digest.update(srcData, 0, srcData.length);
        byte[] hash = new byte[digest.getDigestSize()];
        digest.doFinal(hash, 0);
        return hash;
    }
}
